##########################################################################################

library(data.table)
library(optparse)
library(sva)
library("FactoMineR")
library("factoextra")

##########################################################################################

option_list <- list(
    make_option(c("--expression_file"), type = "character"),
    make_option(c("--base_line_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    expression_file <- "~/20220915_gastric_multiple/rna_batch1/results/RSEM/CombineTPM.tsv"
    base_line_file <- "~/20220915_gastric_multiple/config/Tumor_Normal_RNA.tsv"
    out_path <- "~/20220915_gastric_multiple/rna_batch1/results/compareTpm/"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

expression_file <- opt$expression_file
out_path <- opt$out_path
base_line_file <- opt$base_line_file

dir.create(out_path , recursive = T)

###########################################################################################

dat_expression <- data.frame(fread(expression_file))
info <- data.frame(fread(base_line_file))

###########################################################################################

col_names <- c("ID" , "Sample" , "From" , "Class" )

use_col <- c("ID" , "Tumor" , "From" , "Class" )
dat_tumor <- info[,use_col]
colnames(dat_tumor) <- col_names

use_col <- c("ID" , "Normal" , "From")
dat_normal <- unique(info[,use_col])
dat_normal$Class <- "Normal"
colnames(dat_normal) <- col_names

info_final <- rbind(dat_tumor , dat_normal)

###########################################################################################

colnames(dat_expression) <- gsub( "X" , "" , colnames(dat_expression) )

###########################################################################################
## 看两批次的癌旁表达矩阵是否分开
dat_expression <- dat_expression[,colnames(dat_expression) %in% c( "gene_id" , dat_normal$Sample)]
dat_normal <- subset( dat_normal , Sample %in% colnames(dat_expression) )[,c("Sample" , "From")]

sample_order <- data.frame( Sample = colnames(dat_expression)[-1])
sample_order <- merge( sample_order , dat_normal )


###########################################################################################
## 2022/02/12
## 去除批次效应
## 合并Tumor和Normal
## computeTPM
## https://github.com/ruanjunpeng/eQTLQC/blob/b194e20a0ff128c91ba4cc3968ab0f7dc5397b30/Sample/src/report.Rmd
combatTPM <- function( result_tpm = result_tpm , sample_infor = sample_infor , mod = mod ){
    TPM_data <- result_tpm

    keep_genes_idx <- (rowSums(TPM_data>0.1)>=10) 
    expr= TPM_data[keep_genes_idx,]

    # so row value of 0 will be -4 in the transformed value
    expr=log10(expr+1e-4)  
    # outlier correction: quantile normalization with order preserved. Now TPM is changed to rank normalized gene expression.
    m=apply(expr, 1, mean); 
    sd=apply(expr, 1, sd)
    expr = t(apply(expr, 1, rank, ties.method = "average"));
    #expr = qnorm(expr / (ncol(expr)+1));  # to standard normalization
    expr = qnorm(expr / (ncol(expr)+1), mean=m, sd=sd)  # or, to preserve the mean and sd of each gene
    rm(m,sd)

    expr = as.data.frame(expr)

    if(mod!="NA"){
        combat_tpm <- ComBat(dat = as.matrix(expr) , batch = sample_infor$batch , mod = mod , par.prior=T , prior.plots=FALSE)
    }else{
        combat_tpm <- ComBat(dat = as.matrix(expr) , batch = sample_infor$batch , par.prior=T , prior.plots=FALSE)
    }
    
    combat_tpm <- data.frame(combat_tpm)
    combat_tpm <- data.frame( gene_id = rownames(combat_tpm) , combat_tpm )
    ## 分批次的表达矩阵
    colnames(combat_tpm) <- gsub("[.]" , "-" , colnames(combat_tpm) )

    return(combat_tpm)
}


sample_infor <- sample_order
sample_infor$batch <- sample_infor$From
sample_infor$disease <- "Normal"
mod <- "NA"

result_tpm = dat_expression 
rownames(result_tpm) <- result_tpm$gene_id
result_tpm <- result_tpm[,-1]

## 矫正批次效应
combat_tpm <- combatTPM( result_tpm = result_tpm , sample_infor = sample_infor , mod = mod )


#################################################
## 检查批次效应是否去干净
pca.plot = function(dat , col , out_file){

    df.pca <- PCA(t(dat[,-1]), graph = FALSE)
    fviz_pca_ind(df.pca,
       geom.ind = "point",
       col.ind = col ,
       addEllipses = TRUE,
       legend.title = "Groups"
    )
}

## 原始的批次效应
out_file <- paste0( out_path , "/mRNA_raw_TPM_batch.pdf" )
pdf(out_file)
pca.plot( dat_expression, factor(sample_order$From) , out_file)
dev.off()


## 现在的批次
out_file <- paste0( out_path , "/mRNA_combat_TPM_batch.pdf" )
pdf(out_file)
pca.plot( combat_tpm, factor(paste0(sample_infor$batch , "_" , sample_infor$disease)) , out_file)
dev.off()



